import _pickle
import sys
import cv2
import os
import pickle
import numpy as np
import time

def combinate():
    base = './sample_data/'
    sal_list = ['Conan1', 'Skiing', 'Alien', 'Conan2', 'Surfing', 'War', 'Cooking', 'Football', 'Rhinos']

    x = np.load(base + f'{0}-{sal_list[0]}.npy', allow_pickle=True)
    size = x.shape[0]
    print(size)
    for i in range(1, len(sal_list)):
        y = np.load(base + f'{i}-{sal_list[i]}.npy', allow_pickle=True)
        size += y.shape[0]
        print(y.shape[0])
        x = np.concatenate([x, y])
    np.save('./sample_data/dataset2.npy', x)
    print(x.shape, size)


if __name__ == '__main__':
    combinate()
    exit()

    # videoName = ['1-1-Conan Gore Fly.mp4', '1-2-Front.mp4', '1-3-360 Google Spotlight Stories_ HELP.mp4',
    #             '1-4-Conan Weird Al.mp4', '1-5-TahitiSurf.mp4', '1-6-Falluja.mp4',
    #             '1-7-Cooking Battle.mp4', '1-8-Football.mp4', '1-9-Rhinos.mp4']
    videoNameList = ['Conan1', 'Skiing', 'Alien', 'Conan2', 'Surfing', 'War', 'Cooking', 'Football', 'Rhinos']
    salFileList = ["saliency_ds2_topic0", "saliency_ds2_topic1", "saliency_ds2_topic2",
                   "saliency_ds2_topic3", "saliency_ds2_topic4", "saliency_ds2_topic5",
                   "saliency_ds2_topic6", "saliency_ds2_topic7", "saliency_ds2_topic8"]  # attention数据集
    step = 25

    saliencyBasePath = 'D:/VR_project/PanoSaliency/data/'

    baseId = 1
    total_time = 0
    for i in range(len(videoNameList)):
        start = time.time()
        souredir = videoNameList[i] + '_resize'
        if videoNameList[i] == 'Skiing':
            souredir += '_-20'
        if videoNameList[i] == 'Cooking':
            souredir += '_-19'

        source_path = 'D:/VR_project/saliency-convnet-7.31/framed/' + souredir
        des_path = './sample_frames/'
        saliencyPath = saliencyBasePath + salFileList[i]
        try:
            saliency_array = np.array(pickle.load(open(saliencyPath, 'rb'), encoding='bytes'), dtype=object)
        except _pickle.UnpicklingError:
            saliency_array = np.load(saliencyPath, allow_pickle=True)

        name_list = os.listdir(source_path)
        name_list.sort(key=lambda x: int(x[:-4]))

        idx = 1
        sample_data = []
        while idx * step < len(name_list):
            sample_frame = cv2.imread(source_path + '/' + name_list[idx * step])
            sample_data.append(saliency_array[idx * step - 1])
            cv2.imwrite(des_path + f'{baseId}.jpg', sample_frame)
            idx += 1
            baseId += 1
        end = time.time()
        total_time += round(end - start, 4)
        np.save(f'./sample_data/{i}-{videoNameList[i]}.npy', np.array(sample_data))
        print(f"BaseId={baseId}, sample {len(sample_data)} frames finish for {videoNameList[i]},"
              f" time={round(end - start, 4)}s")

    print(f"Sample {baseId - 1} frames finish in all, time={total_time}s")

